import torch
import numpy as np

out_file = "mbv2_para.h"
bitw, sym_bitw = 8,  "BIT_WT"
bitb, sym_bitb = 24, "BIT_BIAS"
bitm, sym_bitm = 24, "BIT_MULT"

prefixs = [
"conv0",
"b1n1_c1", "b1n1_c2", "b1n1_c3",
"b2n1_c1", "b2n1_c2", "b2n1_c3", "b2n2_c1", "b2n2_c2", "b2n2_c3",
"b3n1_c1", "b3n1_c2", "b3n1_c3", "b3n2_c1", "b3n2_c2", "b3n2_c3", "b3n3_c1", "b3n3_c2", "b3n3_c3",
"b4n1_c1", "b4n1_c2", "b4n1_c3", "b4n2_c1", "b4n2_c2", "b4n2_c3", "b4n3_c1", "b4n3_c2", "b4n3_c3", "b4n4_c1", "b4n4_c2", "b4n4_c3",
"b5n1_c1", "b5n1_c2", "b5n1_c3", "b5n2_c1", "b5n2_c2", "b5n2_c3", "b5n3_c1", "b5n3_c2", "b5n3_c3",
"b6n1_c1", "b6n1_c2", "b6n1_c3", "b6n2_c1", "b6n2_c2", "b6n2_c3", "b6n3_c1", "b6n3_c2", "b6n3_c3",
"b7n1_c1", "b7n1_c2", "b7n1_c3",
"conv8", "conv9"]

def int2hex(x):
	print('\"{:#x}\"'.format(x))
	return '\"{:#x}\"'.format(x)

def write_const(file, symbol, val):
	code = "constexpr unsigned {} = {:d};\n".format(symbol, val)
	file.write(code)
	return

def write_line(file, data, bit):
	assert len(data.shape) == 2
	data[data < 0] += np.power(2, bit)
	outer, inner = data.shape
	line = [None] * outer
	file.write("\t")
	for i in range(outer):
		cat = 0
		for j in range(inner):
			cat |= int(data[i, j]) << (bit * j)
		line[i] = cat
	file.write(", ".join(map(int2hex, line)))
	file.write(",\n")
	return

def write_arr2d(file, data, bit, prefix, sym_width, sym_dim1, sym_dim2):
	assert len(data.shape) == 3
	title = "const ap_uint<{0}> {1}[{2}][{3}] = {{\n".format(sym_width, prefix, sym_dim1, sym_dim2)
	file.write(title)
	for i in range(data.shape[0]):
		write_line(file, data[i], bit)
	file.write("};\n")
	return

def write_arr3d(file, data, bit, prefix, sym_width, sym_dim1, sym_dim2, sym_dim3):
	assert len(data.shape) == 4
	title = "const ap_uint<{0}> {1}_wt[{2}][{3}][{4}] = {{\n".format(sym_width, prefix, sym_dim1, sym_dim2, sym_dim3)
	file.write(title)
	for i in range(data.shape[0]):
		for j in range(data.shape[1]):
			write_line(file, data[i][j], bit)
	file.write("};\n")
	return

def write_bias(file, prefix, n_act):
	bias = np.load("./blob/"+prefix+".bs.npy").astype(np.int32)

def write_dwt(file, prefix, n_ch, n_io):
	wt = np.load("./blob/"+prefix+".wt.npy").astype(np.int32)

def write_pwt(file, prefix, n_i, n_o):
	wt = np.load("./blob/"+prefix+".wt.npy").astype(np.int32)
	n_och = wt.shape[0]
	n_ich = wt.shape[1]
	k = wt.shape[2] * wt.shape[3]

	assert k == 1

	wt_t = wt.reshape(n_och//n_o, n_o, n_ich//n_i, n_i)
	wt_t = wt_t.transpose(0, 2, 1, 3) #(n_och//n_o, n_ich//n_i, n_o, n_i)
	wt_t = wt_t.reshape(n_och//n_o, n_ich//n_i, n_o*n_i)

	sym_ich = "{}_NICH".format(prefix)
	sym_och = "{}_NOCH".format(prefix)
	sym_in  =  "{}_NIN".format(prefix)
	sym_out = "{}_NOUT".format(prefix)
	sym_act = "{}_NACT".format(prefix)
	write_const(file, sym_ich, n_ich)
	write_const(file, sym_och, n_och)
	write_const(file, sym_in , n_i)
	write_const(file, sym_out, n_o)

	sym_width = "{} * {} * {}".format(sym_out, sym_in, sym_bitw)
	sym_dim1 = "{} / {}".format(sym_och, sym_out)
	sym_dim2 = "{} / {}".format(sym_ich, sym_in)

	write_arr2d(file, wt_t, bitw, prefix, sym_width, sym_dim1, sym_dim2)

def write_wt(file, prefix, n_i, n_o):
	wt = np.load("./blob/"+prefix+".wt.npy").astype(np.int32)
	n_och = wt.shape[0]
	n_ich = wt.shape[1]
	k = wt.shape[2] * wt.shape[3]
	wt_t = wt.reshape(n_och//n_o, n_o, n_ich//n_i, n_i, k)
	wt_t = wt_t.transpose(0, 2, 4, 1, 3) #(n_och//n_o, n_ich//n_i, k, n_o, n_i)
	wt_t = wt_t.reshape(n_och//n_o, n_ich//n_i, k, n_o*n_i)

	sym_ich = "{}_NICH".format(prefix)
	sym_och = "{}_NOCH".format(prefix)
	sym_in  =  "{}_NIN".format(prefix)
	sym_out = "{}_NOUT".format(prefix)
	sym_k   = "{}_K".format(prefix, prefix)
	sym_act = "{}_NACT".format(prefix)
	write_const(file, sym_ich, n_ich)
	write_const(file, sym_och, n_och)
	write_const(file, sym_in , n_i)
	write_const(file, sym_out, n_o)
	write_const(file, sym_k, wt.shape[2])
	sym_width = "{} * {} * {}".format(sym_out, sym_in, sym_bitw)
	sym_dim1 = "{} / {}".format(sym_och, sym_out)
	sym_dim2 = "{} / {}".format(sym_ich, sym_in)
	sym_dim3 = "{} * {}".format(sym_k, sym_k)
	write_arr3d(file, wt_t, bitw, prefix, sym_width, sym_dim1, sym_dim2, sym_dim3)

f = open(out_file, "w")
write_wt(f, prefixs[0], 3, 4)
#write_pwt(f, prefixs[1], 4, 4)